import datetime
import easyquotation
import pandas as pd 
import os
import numpy as np 
import warnings
warnings.filterwarnings('ignore')
import sys
sys.path.append(r'code')
from real.email_test import send_email
import pickle

#阈值加rank开仓方式，附带止盈止损
class Pos_make_rank_thred_open():   
    
    def __init__(self, args, open_root=None, close_root=None):
        self.args = args
        self.pos_root = self.pos_root_make()
        self.today = pd.to_datetime(datetime.date.today())     
        self.money_init = args.money_init
        self.thred = args.open_thred
        self.predict = self.get_predict()   
        self.open, self.close = self.get_price(open_root, close_root)
        self.hold_ticker = self.hold_ticker_init() # 记录持仓量价
        self.pos_init()
        self.cost_ratio = args.cost_ratio
        self.hold_num = args.hold_num
    
    def log_write(self, content):
        f = os.path.join(self.pos_root, 'log.txt')
        with open(f,"a") as file:
            file.write(content)    
    
    def hold_ticker_init(self):
        file = os.path.join(self.pos_root, 'ticker.pkl')
        if os.path.exists(file):
            with open(file, 'rb') as f:
                data = pickle.load(f)
        else:
            data = {}
        return data
           
    def pos_init(self):
        self.buy_pos = self.read_pos('buy_pos')
        self.sell_pos = self.read_pos('sell_pos')
        self.hold_pos = self.read_pos('hold_pos')
        self.money_df = self.read_money()
        self.cost_df = self.read_pos('cost_df')
        self.money = self.money_df.iloc[-1].values[0]
        self.log_write( f'\nthe last day, buy_pos:{self.buy_pos.index[-1]}, sell_pos:{self.sell_pos.index[-1]}, hold_pos{self.hold_pos.index[-1]}, money_df:{self.money_df.index[-1]}' )
        print( f'the last day, buy_pos:{self.buy_pos.index[-1]}, sell_pos:{self.sell_pos.index[-1]}, hold_pos{self.hold_pos.index[-1]}, money_df:{self.money_df.index[-1]}' )

    def read_pos(self, name='buy_pose'):
        pos_file = os.path.join(self.pos_root, f'{name}.csv')
        contract = self.predict.columns
        if os.path.exists(pos_file):
            pos_df = pd.read_csv(pos_file, index_col=0)
        else:
            zero_ary = np.zeros((len(self.predict),len(contract)))
            pos_df = pd.DataFrame(zero_ary, index= self.predict.index, columns=contract)
        pos_df.index = pd.to_datetime(pos_df.index)
        pos_df = pos_df.astype(int)
        
        pos_today = pd.DataFrame(np.zeros((1,len(contract))), index = [self.today], columns=contract)
        pos_df = pd.concat([pos_df, pos_today], axis=0)
        pos_df = pos_df.reset_index().drop_duplicates(subset='index', keep='first').set_index('index')
        pos_df = pos_df.fillna(0)
        pos_df = pos_df[sorted(pos_df.columns)]
        # pos_df.loc[self.today] = np.zeros(len(contract))
        return pos_df

    def read_money(self):
        money_file = os.path.join(self.pos_root, 'money.csv')
        if os.path.exists(money_file):
            money_df = pd.read_csv(money_file, index_col=0)
        else:
            zero_ary = np.ones(len(self.predict))*self.money_init
            money_df = pd.DataFrame(zero_ary, index= self.predict.index, columns=['money'])
        money_df.index = pd.to_datetime(money_df.index)
        money_df.loc[self.today] = np.nan
        money_df.loc[self.today] = money_df.values[-2]
        return money_df

    def pos_root_make(self):
        result_root = os.path.join(self.args.result_root, 'pos')
        if not os.path.exists(result_root):
            os.makedirs(result_root)
        return result_root

    def get_price(self, open_root, close_root):  
        # sourcery skip: avoid-builtin-shadow
        if open_root is None:
            open_root = r'data/stock_data/consentrate_daily_price/open.pkl.gzip'
        if close_root is None:
            close_root = r'data/stock_data/consentrate_daily_price/close.pkl.gzip'
        open = pd.read_pickle(open_root)
        close = pd.read_pickle(close_root)

        # 新浪 ['sina'] 腾讯 ['tencent', 'qq']
        quotation = easyquotation.use('tencent')
        tmp_dict = quotation.market_snapshot(prefix=True)
        today = str(datetime.date.today())
        for code in open.columns:
            key = f'{code[:2]}{code[3:]}'
            close.loc[today, code] = tmp_dict[key]['close']
            open.loc[today, code] = tmp_dict[key]['now']
        close.index = pd.to_datetime(close.index)
        open.index = pd.to_datetime(open.index)
        close = close.iloc[-10:]
        open = open.iloc[-10:]
        return open, close

    def get_predict(self):
        df = pd.read_pickle(os.path.join(self.args.result_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.loc[self.today] = np.zeros(len(df.columns))
        df = df.shift(1) # 预测值滞后一天
        return df
    
    def stop_ratio(self): #计算平均成本
        idv = self.today
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        sell_dict = {}
        for code in list(self.hold_ticker.keys()):          
            price_cost = self.hold_ticker[code][1]
            price_now = self.open.loc[idv, code]
            ratio = 0
            ratio = price_now/price_cost -1
            #止盈平仓, #止损平仓  
            if (ratio >= self.stop_ret) or (ratio <= self.stop_loss):
                self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
                hold_num = self.hold_ticker[code][0]
                if hold_num > 0: #持仓大于0
                    self.sell_pos.loc[idv, code] = hold_num
                    price = self.open.loc[idv, code]
                    self.money += (price * 100 *hold_num)
                    self.log_write(f'\n{idv}：{code}触发止损止盈平仓 , 剩余金额：{self.money:.2f},  收益率：{ratio:.2f}')
                    print(f'{idv}：{code}触发止损止盈平仓 ,  剩余金额：{self.money:.2f}, 收益率：{ratio:.2f}')
                    self.hold_ticker.pop(code)
                    cost = np.maximum(price * 100 *hold_num*self.cost_ratio, 1)
                    self.cost_df.loc[idv, code] += cost
                    sell_dict[code] = hold_num
        return sell_dict
    
    def pos_add(self):
        idv = self.today
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        price_limit = 100
        pred = self.predict.loc[idv]
        open_v = self.open.loc[idv]
        open_v = open_v[(open_v >= 5) & (open_v < price_limit) ]
        pred = pred[open_v.index]
        pred = pred[pred>=0.6]
        pred = pred.sort_values(ascending=False)
        pred = pred[:self.hold_num]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)

        while (self.money > cost_lst+price_list).any():
            for code in buy_codes:#debug 检查此处             
                price_amount = self.open.loc[idv, code] * 100
                cost = np.maximum(price_amount*self.cost_ratio, 1)
                if self.money >= (price_amount+cost): #资金满足开仓要求
                    buy_dict[code] += 1
                    price = self.open.loc[idv, code]
                    self.money -= (price * 100) 
                    self.cost_df.loc[idv, code] += cost  
                    self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                    print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = num
            if num > 0:
                if code not in list(self.hold_ticker.keys()):
                    self.hold_ticker[code] = [num, self.open.loc[idv, code]]
                else:
                    cnum = self.hold_ticker[code][0] + num
                    pric = (self.hold_ticker[code][0]*self.hold_ticker[code][1]+num*self.open.loc[idv, code])/(self.hold_ticker[code][0]+num)
                    self.hold_ticker[code] = [cnum, pric]
            else:
                buy_dict.pop(code)
        return buy_dict

    def pos_make(self): 
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()  
        sell_dict =  self.stop_ratio()
        buy_dict = self.pos_add() 
        buy_dict['direct'] = 'cls buy'
        sell_dict['direct'] = 'cls sell'     
        if self.args.send_email:
            send_email([sell_dict, buy_dict])
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        self.pos_result_save()                           
                
    def pos_result_save(self):
        self.buy_pos.to_csv(f'{self.pos_root}/buy_pos.csv')
        self.sell_pos.to_csv(f'{self.pos_root}/sell_pos.csv')
        self.hold_pos.to_csv(f'{self.pos_root}/hold_pos.csv')
        self.money_df.to_csv(f'{self.pos_root}/money.csv')
        self.cost_df.to_csv(f'{self.pos_root}/cost.csv')
        with open(f'{self.pos_root}/ticker.pkl', 'wb') as f:
            pickle.dump(self.hold_ticker, f)